#include "actions/ferm/invert/invcg2_array.h"
#include "fermact.h"
#include <actions/ferm/invert/syssolver_linop_cg_array.h>
#define gaugeSiteSize 18 // real numbers per link
#include "chroma.h"
using namespace Chroma;
#include "dwf_4d_invert.h"
#include <complex>
#include "reunit.h"

enum { DaggerNo, DaggerYes };

void qdp_dwf_invert( multi1d<LatticeFermion> &eta, multi1d<LatticeFermion> &psi, multi1d<LatticeColorMatrix> &u,double M5);

/*
 * Test application that will make use of the 4d_dwf_invert class
 */
int main (int argc,char **argv )
{
    Chroma::initialize(&argc, &argv);

    /*
     * Read commandline arguments for test program
     */
    if ( argc != 6 )
    {
        printf("Usage: %s Lx Ly Lz Lt Ls\n All must be even\n",argv[0]);
        exit(-1);
    }

    /*
     * Setting up the QDP enviroment
     */
    multi1d<int> nrow(Nd);
    nrow[0] = atoi(argv[1]);
    nrow[1] = atoi(argv[2]);
    nrow[2] = atoi(argv[3]);
    nrow[3] = atoi(argv[4]);
    int Ls = atoi(argv[5]);

    Layout::setLattSize(nrow);
    Layout::create();

    /*
     * Create our '5th' dimension as array of fermions
     */
    multi1d<LatticeFermion> eta(Ls);
    multi1d<LatticeFermion> fteta(Ls);
    multi1d<LatticeFermion> psi(Ls);
    multi1d<LatticeFermion> iftpsi(Ls);
    multi1d<LatticeFermion> qdppsi(Ls);

    //gaussian(eta[0]);
    for ( int s = 0; s < (Ls); s++ )
    {
        eta[s] = zero;
        //gaussian(eta[s]);
        //eta[s].elem(all.start()).elem(0).elem(2).real() = s;
        //eta[s].elem(all.start()).elem(0).elem(8).real() = Ls - s;
    }
    gaussian(eta[0]);
    gaussian(eta[Ls-1]);
   // eta[Ls-1].elem(all.start()).elem(0).elem(0).real() = 1;
   // eta[0].elem(all.start()).elem(0).elem(17).real() = 1;

    multi1d<LatticeColorMatrix> u(Nd);
    for ( int mu = 0; mu < Nd; mu++ )
    {
        u[mu] = zero;
        // gaussian(u[mu]);
        //TEST_LIB::reunit(u[mu]);
       //     u[mu] = Real(1.0);
        HotSt(u);
    }

    initQuda(0);

    QDPIO::cout << "////////////////////////////////////////////\n";
    QDPIO::cout << "// Finding result using FT and QUDA         \n";
    QDPIO::cout << "////////////////////////////////////////////\n";

    QDPIO::cout << "\n~Source Norm2 = " << norm2(eta) << " \n\n";
    QDPIO::cout << "Fourier Transforming Source\n";
    dwf_4d_invert::fourierTransform(eta, fteta);
    QDPIO::cout << "~FT Source Norm2 = " << norm2(fteta) << " \n\n";


    QDPIO::cout << "Solving for the solution vector\n";
    dwf_4d_invert::domainWallFermionInvert( fteta, psi, u, 1.8, DaggerNo);
    QDPIO::cout << "~Solution Norm2 = " << norm2(psi) << " \n\n";


    QDPIO::cout << "Inverse Fourier Transforming Solution\n";
    dwf_4d_invert::fourierTransformInv(psi, iftpsi);
    QDPIO::cout << "~Inverse FT Solution Norm2 = " << norm2(iftpsi) << " \n\n";

    QDPIO::cout << "////////////////////////////////////////////\n";
    QDPIO::cout << "// Finding result using QDP                 \n";
    QDPIO::cout << "////////////////////////////////////////////\n";

    QDPIO::cout << "\n~Source Norm2 = " << norm2(eta) << " \n\n";

    QDPIO::cout << "Computing result using QDP direct invert\n";
    qdp_dwf_invert(eta, iftpsi, u, 1.8);
    QDPIO::cout << "~QDP Solution Norm2 = " << norm2(qdppsi) << " \n\n";

    QDPIO::cout << "Comparing Results\n";
    QDPIO::cout << "~4dSolution - QDPSolution Norm2 = " << norm2(iftpsi-qdppsi) << " \n";


    endQuda();

    Chroma::finalize();
}




void qdp_dwf_invert( multi1d<LatticeFermion> &eta, multi1d<LatticeFermion> &psi, multi1d<LatticeColorMatrix> &u,double m5){
    /********************************************************
     * QDP Linop
     ********************************************************
     */
    typedef LatticeFermion T;
    typedef multi1d<LatticeFermion> T5;
    typedef multi1d<LatticeColorMatrix> U;

    int Ls = eta.size();

    multi1d<LatticeFermion> qdp_eta(eta.size());

    Real M5(m5);
    Real mass = 1.00;
    Real mq = mass;
#define Ncg 3
    Real residuals[Ncg] ; 
    int  max_iter[Ncg];
    residuals[0] = 1.0e-6;
    residuals[1] = 1.0e-7;
    residuals[2] = 1.0e-8;
    //  residuals[3] = 1.0e-8;

    for(int i=0;i<Ncg;i++) max_iter[i]  = 1000;


    multi1d<int> bcs(Nd);
    bcs[0] = bcs[1] = bcs[2] = bcs[3] = 1;

    Handle< FermBC<T,U,U> > fbc(new SimpleFermBC< T, U, U >(bcs));
    Handle<CreateFermState<T,U,U> > cfs( new CreateSimpleFermState<T,U,U>(fbc));
    UnprecDWFermActArray  S_f(cfs, M5, mq, Ls);
    Handle< FermState<T,U,U> > fs( S_f.createState(u) );

    Handle< UnprecLinearOperatorArray<T,U,U> > M(S_f.unprecLinOp(fs,mq));

    (*M)(qdp_eta,psi,PLUS);

    QDPIO::cout << "s | norm2(eta[s])    | norm2(psi[s])    | norm2(qdp_eta[s])| norm2(qdp_eta[s]-eta[s])| norm2(eta[s])/norm2(qdp_eta[s])" << endl;
    for ( int i = 0; i < eta.size(); i++ )
    {
        QDPIO::cout << i << " | " << norm2(eta[i]) << " | " << norm2(psi[i]) << " | " << norm2(qdp_eta[i]) << " | " << norm2(qdp_eta[i] - eta[i]) << " | " << norm2(eta[i])/norm2(qdp_eta[i]) << endl;
    }
}
